import serial

def to_labbrick_readable_freq(freq, mode=0):
    if mode==0:
        out = "F"+format(format(freq,".2f"),"0>7s")[-7:]
    if mode==1:
        out = "F"+format(format(freq,".2f"),"0>8s")[-8:]
    return out

def to_labbrick_readable_amp(amp):
    return "A"+format(format(amp,".0f"),"0>2s")[-2:]+"ATTEN"

class Labbrick:
    def __init__(self, serial_port = None, mode=0):
        self.port = serial_port
        self.mode = mode # 0 represents ordinary mode; 1 represents high frequency mode (used for generate 13GHz)
        
        if serial_port != None:
            self.ser  = serial.Serial(
                port=serial_port,
                baudrate=9600,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_TWO,
                bytesize=serial.EIGHTBITS
            )
            
    def freq_update(self, freq):
        str_to_send = to_labbrick_readable_freq(freq, self.mode).encode('utf_8')
        self.ser.write(str_to_send)
        
    def amp_update(self, amp):
        if(amp > 0 or amp < -30):
            raise AssertionError("The amp of domestic Labbrick should in [-30, 0], which represents the attenuation of the output. Check Feishu for detail.")
        else:
            str_to_send = to_labbrick_readable_amp(abs(amp)).encode('utf_8')
            self.ser.write(str_to_send)